import math, sys

g_RARE = "_RARE_"

def read_counts_file(filename):
    counts = {}
    f = open(filename)
    for line in f:
        tokens = line.split()
        if len(tokens) <> 4:
            continue
        if tokens[1] == "WORDTAG":
            count = int(tokens[0])
            tag = tokens[2]
            word = tokens[3]
            counts[(tag, word)] = count
    f.close()
    return counts

def compute_prob(counts):
    counts = collapse_rare_words(counts)
    prob = {}
    tag_counts = {}
    for pair, count in counts.viewitems():
        tag = pair[0]
        accum(tag_counts, tag, count)
    for pair, count in counts.viewitems():
        tag = pair[0]
        word = pair[1]
        tag_count = tag_counts[tag]
        dict2_add(prob, word, tag, count * 1.0 / tag_count)
    return prob

def collapse_rare_words(counts):
    word_counts = {}
    for pair, count in counts.viewitems():
        word = pair[1]
        accum(word_counts, word, count)
    rare_words = set([word for word, count in word_counts.viewitems() if count < 5])
    new_counts = {}
    for pair, count in counts.viewitems():
        tag = pair[0]
        word = pair[1]
        if word in rare_words:
            word = g_RARE
        accum(new_counts, (tag, word), count)
    return new_counts

def accum(d, key, value):
    d[key] = d.get(key, 0) + value

def dict2_add(d, key1, key2, value):
    d2 = d.get(key1, {})
    d2[key2] = value
    d[key1] = d2

def tag_file(sent_file, prob):
    f = open(sent_file)
    for line in f:
        tokens = line.split()
        if len(tokens) <> 1:
            print
            continue
        word = tokens[0]
        tag_prob = prob.get(word, None)
        if tag_prob == None:
            tag_prob = prob[g_RARE]
        tag, p = max(tag_prob.viewitems(), key = lambda x : x[1])
        print word, tag, math.log(p, 2)
    f.close

counts_file = "ner.counts"
sent_file = "ner_dev.dat"

if len(sys.argv) >= 3:
    counts_file = sys.argv[1]
    sent_file = sys.argv[2]

counts = read_counts_file(counts_file)
# print counts
prob = compute_prob(counts)
# print prob
tag_file(sent_file, prob)
